#该问题发生于当我使用yolo官方模型推理时，其生成的txt文件有我不需要的类和类发生错位
#所以需要更改生成txt文件的内容
import argparse
from  pathlib import Path
import numpy as np
#定义可由命令行输入的参数
def Parse_Arguments():
    parser = argparse.ArgumentParser(description="对文件夹及文件处理")
    #源路径
    parser.add_argument("--origin_dir_path", type=str, default="D:/tempDownload/dataset/elevator/train/img1/labels")
    return parser.parse_args()

def ProcessTxt(source_label_content):
    ans = []
    if np.array(source_label_content).shape[0] == 0:
        return ans

    if len(np.array(source_label_content).shape) == 1:
        if source_label_content[0] == 16:
            source_label_content[0] = 1
            ans = source_label_content
        return ans

    for item in source_label_content:
        if item[0] != 16 and item[0] != 0:
            continue
        if item[0] == 16:
            item[0] = 1
        ans.append(list(item))
    return ans

def main():
    args = Parse_Arguments()
    source_label_files = Path(args.origin_dir_path).glob("*_*.txt")
    for source_label_file in source_label_files:
        try:
            source_label_content = np.loadtxt(str(source_label_file))
            target_laebl_content = ProcessTxt(source_label_content)
            # 定义数据格式
            fmt = ['%.0f'] + ['%f'] * 4
            np.savetxt(str(source_label_file), target_laebl_content, fmt=fmt)
        except:
            print(f"{source_label_file}出现问题")
            continue
if __name__ == "__main__":
    main()
